import os
import shutil

scp_out_base_dir = "/share/home/wuqingyao_danglingwei/datas/TACO_Data_20250314/test_datas"

data_root = "/share/home/wuqingyao_danglingwei/datas/TACO_Data_20250314/full_20k_plus"
prompts_column_path = "test_small_151_VLMEnhanced_prompts.txt"
img_column_path = "test_small_151_images.txt"
color_video_column_path = "test_small_151_color_videos.txt"
pose_video_column_path = "test_small_151_black_pose_videos.txt"
hoi_traj_path = "test_small_151_normalized_hoi_pose_trajs.txt"

def get_file_path_list(column_path, mode="prompts"):
    result_list = []
    cnt = 0
    with open(column_path, "r", encoding="utf-8") as f:
        for line in f:
            line = line.strip()
            if line:
            # if line and cnt < total_gen_sample:
                cnt+=1
                if mode == "prompts":
                    result_list.append(line)
                else:
                    result_list.append(line)
            else:
                break
    return result_list
    
prompts_list = get_file_path_list(os.path.join(data_root, prompts_column_path), mode="prompts")
img_path_list = get_file_path_list(os.path.join(data_root, img_column_path), mode="others")
color_video_path_list = get_file_path_list(os.path.join(data_root, color_video_column_path), mode="others")
pose_video_path_list = get_file_path_list(os.path.join(data_root, pose_video_column_path), mode="others")
hoi_pose_traj_path_list = get_file_path_list(os.path.join(data_root, hoi_traj_path), mode="others")

assert len(prompts_list) == len(img_path_list) == len(color_video_path_list) == len(pose_video_path_list) == len(hoi_pose_traj_path_list), "The length of the lists do not match!"

for idx, (prompt, img_path, color_video_path, pose_video_path, hoi_pose_traj_path) in enumerate(zip(prompts_list, img_path_list, color_video_path_list, pose_video_path_list, hoi_pose_traj_path_list)):
    print(f"idx: {idx}")
    from_img_path = os.path.join(data_root, img_path)
    to_img_path = os.path.join(scp_out_base_dir, img_path)

    # print(f"from_img_path: {from_img_path}, to_img_path: {to_img_path}")
    os.makedirs(os.path.dirname(to_img_path), exist_ok=True)
    shutil.copyfile(from_img_path, to_img_path)

    from_color_video_path = os.path.join(data_root, color_video_path)
    to_color_video_path = os.path.join(scp_out_base_dir, color_video_path)
    os.makedirs(os.path.dirname(to_color_video_path), exist_ok=True)
    shutil.copyfile(from_color_video_path, to_color_video_path)

    from_pose_video_path = os.path.join(data_root, pose_video_path)
    to_pose_video_path = os.path.join(scp_out_base_dir, pose_video_path)
    os.makedirs(os.path.dirname(to_pose_video_path), exist_ok=True)
    shutil.copyfile(from_pose_video_path, to_pose_video_path)

    from_pose_traj_path = os.path.join(data_root, hoi_pose_traj_path)
    to_pose_traj_path = os.path.join(scp_out_base_dir, hoi_pose_traj_path)
    os.makedirs(os.path.dirname(to_pose_traj_path), exist_ok=True)
    shutil.copyfile(from_pose_traj_path, to_pose_traj_path)

print(f"Done!")
